# -*- coding: UTF-8 -*-

# %% [markdown]
import os
import re
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import sys
import sympy as sp


def fourier_series(coeffs):
    theta = sp.symbols('theta')
    return np.sum([coeff[0]*sp.cos(coeff[1]*theta+coeff[2]) for coeff in coeffs])


def my_fft(Ts, y: np.ndarray):
    sp = np.fft.fft(y)
    freq = np.fft.fftfreq(y.shape[-1])
    amplitude = (abs(sp)/len(y)*2)[:int(len(freq)/2)]
    freq_Hz = (freq/Ts)[:int(len(freq)/2)]
    phase = np.arctan(abs(sp.imag)/abs(sp.real))[:int(len(freq)/2)]

    return freq_Hz, amplitude, phase


if __name__ == "__main__":
    pass

    # %% [markdown]
    # 模拟真实编码器数据
    encoder_max_pulse = 16384
    rad2encoder = encoder_max_pulse/(2*np.pi)
    encoder2rad = (2*np.pi)/encoder_max_pulse

    num = np.arange(0, 200)
    theta_idel = num/200*2*np.pi
    # plt.scatter(num, theta_idel)

    theta_ripple = 0.2*np.cos(2*theta_idel)
    # + \
    #     0.08*np.cos(4*theta_idel + np.pi/2)
    #     0.2*np.cos(8*theta_idel) + \
    #     0.1*np.cos(14*theta_idel)

    theta_real = theta_idel - theta_ripple

    encoder_real = theta_real / (2*np.pi)*2**14

    theta_err = theta_idel - theta_real
    plt.figure(0)
    plt.plot(num, theta_ripple, marker='.')
    plt.figure(1)
    plt.plot(theta_idel, theta_err, marker='.')
    plt.figure(2)
    plt.plot(theta_real, theta_err, marker='.')
    # plt.plot(num, encoder_real, marker='.')

    plt.show()

    # %% [markdown]
    # 读取真实数据
    with open('./data.txt', 'r') as f:
        line = f.readline().splitlines()

        i = 0
        data_list = []
        while line:
            i += 1
            data = str(*line).split(' ')
            data = [float(x) for x in data]
            data_list.append(data)
            # print ("{}: {} | {}".format(i, data, type(data)))
            line = f.readline().splitlines()
        print("size:{}".format(len(data_list)))
    data_list = np.array(data_list)
    theta_idel = np.array(data_list[:, 0]*encoder2rad)
    theta_real = np.array(data_list[:, 1]*encoder2rad)
    theta_err = theta_idel - theta_real
    plt.figure(0)
    plt.plot(theta_idel, marker='.')
    plt.figure(1)
    plt.plot(theta_real, marker='.')
    plt.figure(2)
    plt.plot(theta_err, marker='.')
    plt.figure(3)
    plt.plot(theta_real, theta_err, marker='.')
    # %% [markdown]
    # 拟合

    def func8(theta, c, a1, phi1, 
                            a2, phi2, 
                            a3, phi3, 
                            a4, phi4,
                            a5, phi5, 
                            a6, phi6, 
                            a7, phi7, 
                            a8, phi8):
        return a1 * np.cos(theta + phi1) + \
                a2 * np.cos(2 * theta + phi2) + \
                a3 * np.cos(3 * theta + phi3) + \
                a4 * np.cos(4 * theta + phi4) + \
                a5 * np.cos(5 * theta + phi5) + \
                a6 * np.cos(6 * theta + phi6) + \
                a7 * np.cos(7 * theta + phi7) + \
                a8 * np.cos(8 * theta + phi8) + \
                c
    popt,pcov=curve_fit(func8, theta_real, theta_err)
    plt.figure(0)
    plt.plot(theta_real, func8(theta_real, *popt))
    plt.plot(theta_real, theta_err)
    plt.figure(1)
    
    # plt.plot(theta_idel-(theta_real + theta_err))
    plt.plot(theta_idel-(theta_real + func8(theta_real, *popt)))
    plt.plot(theta_idel-(theta_real + theta_err))
    for i in range(0,len(popt[1:]),2):
        print('{{{}, {}, {}}},'.format(popt[1+i], i/2+1, popt[1+i+1]))